package org.rbio.paper.methylglm.glm;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileReader;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.Collections;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.TreeMap;

import org.apache.commons.math3.distribution.ChiSquaredDistribution;
import org.kohsuke.args4j.CmdLineParser;
import org.kohsuke.args4j.Option;

// parse the config file, merge the data files and evaluate GLM deviance on the table
// config file has the following format:
// [expt. condition]	[file name with path]
// data files must have the follwoing format:
// [chr]	[pos]	[meth_read]	[not_meth_read]
public class MethylGLMFit {
	// defaults
	public static final int DEFAULT_INT = -99999;

	// command line arguments
	@Option(name = "-config", usage = "configuration file", required = true)
	private String config_file;

	@Option(name = "-min_tot_read", usage = "minimum number of total reads per CpG", required = true)
	private int min_tot_read;

	@Option(name = "-out_prefix", usage = "prefix for output files", required = true)
	private String out_prefix;

	// members
	Config config;
	Map<CpG, List<DataPoint>> data_tab;

	// entry point
	public static void main(String[] args) throws Exception {
		MethylGLMFit prep = new MethylGLMFit();
		CmdLineParser parser = new CmdLineParser(prep);
		try {
			parser.parseArgument(args);
			prep.process();
		} catch (Exception e) {
			System.err.println(e.getMessage());
			System.err
					.println("Perform GLM on a single CpG across the different experimental conditions and the replicates");
			parser.printUsage(System.err);
			return;
		}
	}

	// process
	public void process() throws Exception {
		try {
			// parse config file
			config = new Config(config_file);

			// merge the files and sort the CpGs by position
			merge();

			// evaluate GLM deviance
			calculateDeviance();

			// make the FDR table
			calculateFDR();

			// write the results
			write();
		} catch (Exception e) {
			e.printStackTrace();
		}
	}

	// parse the config and merge the data files
	@SuppressWarnings("static-access")
	public void merge() throws Exception {
		data_tab = new TreeMap<CpG, List<DataPoint>>();
		for (String expt : config.expt_tab.keySet()) {
			for (Replicate rep : config.expt_tab.get(expt)) {
				BufferedReader ibuf = new BufferedReader(new FileReader(
						rep.file));
				String line = ibuf.readLine();
				while ((line = ibuf.readLine()) != null) {
					// check for empty lines
					if (line.length() == 0) {
						continue;
					}

					// add data line to the table
					String arr[] = line.split("\t");
					int meth_read = Integer.parseInt(arr[rep.METH_READ_COL]);
					int not_meth_read = Integer
							.parseInt(arr[rep.NOT_METH_READ_COL]);

					// check if total number of reads is enough [user specified]
					if (meth_read + not_meth_read < min_tot_read) {
						continue;
					}

					// add data for this CpG
					CpG cpg = new CpG(arr[rep.CHR_COL],
							Integer.parseInt(arr[rep.POS_COL]));
					List<DataPoint> rec_lst = data_tab.get(cpg);
					if (rec_lst == null) {
						rec_lst = new ArrayList<DataPoint>();
						data_tab.put(cpg, rec_lst);
					}

					// check
					rec_lst.add(new DataPoint(meth_read, not_meth_read));
				}
				ibuf.close();
			}
		}

		// filter out the CpG positions where data is missing for some of the
		// replicates
		List<CpG> cpg_lst = new ArrayList<CpG>();
		for (CpG cpg : data_tab.keySet()) {
			List<DataPoint> rec_lst = data_tab.get(cpg);
			if (rec_lst.size() != config.num_rep) {
				cpg_lst.add(cpg);
			}
		}
		for (CpG cpg : cpg_lst) {
			data_tab.remove(cpg);
		}
	}

	// deviance calculation
	public void calculateDeviance() {
		// find deviance for each CpG
		for (CpG cpg : data_tab.keySet()) {
			List<DataPoint> rec_lst = data_tab.get(cpg);
			int tot_meth = 0, tot_not_meth = 0;
			List<Integer> expt_meth_lst = new ArrayList<Integer>();
			List<Integer> expt_not_meth_lst = new ArrayList<Integer>();
			int idx = 0;
			for (String expt : config.expt_tab.keySet()) {
				int expt_meth = 0, expt_not_meth = 0;
				for (int i = 0; i < config.expt_tab.get(expt).size(); i++) {
					expt_meth += rec_lst.get(idx).meth_read;
					expt_not_meth += rec_lst.get(idx).not_meth_read;
					idx++;
				}
				expt_meth_lst.add(expt_meth);
				expt_not_meth_lst.add(expt_not_meth);

				// update total
				tot_meth += expt_meth;
				tot_not_meth += expt_not_meth;
			}

			// find deviance for this CpG
			double expt_dev = deviance(expt_meth_lst, expt_not_meth_lst,
					rec_lst, false);
			double tot_dev = deviance(tot_meth, tot_not_meth, rec_lst);
			cpg.dev = tot_dev - expt_dev;
		}
	}

	// deciance prep [for the null model]
	public double deviance(int meth, int not_meth, List<DataPoint> rec_lst) {
		List<Integer> meth_lst = new ArrayList<Integer>();
		meth_lst.add(meth);
		List<Integer> not_meth_lst = new ArrayList<Integer>();
		not_meth_lst.add(not_meth);
		return deviance(meth_lst, not_meth_lst, rec_lst, true);
	}

	// deviance [list]
	public double deviance(List<Integer> meth_lst, List<Integer> not_meth_lst,
			List<DataPoint> rec_lst, boolean is_null_model) {
		double sum = 0;
		int expt_idx = 0, rep_idx = 0;
		for (String expt : config.expt_tab.keySet()) {
			int expt_meth = meth_lst.get(expt_idx);
			int expt_not_meth = not_meth_lst.get(expt_idx);
			double expt_meth_est = expt_meth
					/ (expt_meth + expt_not_meth * 1.0);
			for (int i = 0; i < config.expt_tab.get(expt).size(); i++) {
				int rep_meth = rec_lst.get(rep_idx).meth_read;
				int rep_not_meth = rec_lst.get(rep_idx).not_meth_read;
				int rep_tot = rep_meth + rep_not_meth;

				// check if meth = 0
				double prt1 = 0;
				if (rep_meth != 0) {
					prt1 = rep_meth
							* Math.log(rep_meth / (rep_tot * expt_meth_est));
				}

				// check if not_meth = 0
				double prt2 = 0;
				if (rep_not_meth != 0) {
					prt2 = rep_not_meth
							* Math.log(rep_not_meth
									/ (rep_tot * (1 - expt_meth_est)));
				}

				// update sum
				sum += prt1 + prt2;

				// update replicate index
				rep_idx++;
			}

			// update experiment index
			if (is_null_model == false) {
				expt_idx++;
			}
		}
		return 2 * sum;
	}

	// fdr table
	public void calculateFDR() throws Exception {
		// degrees of freedom
		int df = config.expt_tab.size() - 1;
		ChiSquaredDistribution chisq = new ChiSquaredDistribution(df);

		// calculate p-values
		List<PValue> p_lst = new ArrayList<PValue>();
		for (CpG cpg : data_tab.keySet()) {
			p_lst.add(new PValue(cpg.dev, 1 - chisq
					.cumulativeProbability(cpg.dev)));
		}

		// sort the p-values
		Collections.sort(p_lst);

		// fdr for which deviance cutoffs are evaluated
		double fdr_arr[] = { 0.01, 0.05, 0.1, 0.2 };

		// find the deviance cutoffs
		List<Double> dev_cut_lst = new ArrayList<Double>();
		for (Double fdr : fdr_arr) {
			double fdr_dev = DEFAULT_INT;
			for (int i = 0; i < p_lst.size(); i++) {
				PValue pval = p_lst.get(i);
				if (pval.pval > (i + 1.0) / p_lst.size() * fdr) {
					fdr_dev = pval.dev;
					break;
				}
			}
			dev_cut_lst.add(fdr_dev);
		}

		// write
		BufferedWriter obuf = new BufferedWriter(new FileWriter(out_prefix
				+ ".fdr.txt"));
		obuf.write("FDR\tdeviance\n");
		for (int i = 0; i < fdr_arr.length; i++) {
			obuf.write(fdr_arr[i] + "\t" + dev_cut_lst.get(i) + "\n");
		}
		obuf.close();
	}

	// write results
	public void write() throws Exception {
		// write the merged data file
		BufferedWriter obuf = new BufferedWriter(new FileWriter(out_prefix
				+ ".merge.txt"));

		// write the header
		String header = "chr\tpos";
		for (String expt : config.expt_tab.keySet()) {
			for (int i = 1; i <= config.expt_tab.get(expt).size(); i++) {
				String rep_str = "\t" + expt + "_" + i;
				header += rep_str + "_meth" + rep_str + "_not_meth";
			}
		}
		obuf.write(header + "\tdev" + "\n");

		// write the data
		for (CpG cpg : data_tab.keySet()) {
			List<DataPoint> rec_lst = data_tab.get(cpg);

			String res = cpg.chr + "\t" + cpg.pos;
			for (DataPoint rec : rec_lst) {
				res += "\t" + rec.meth_read + "\t" + rec.not_meth_read;
			}
			obuf.write(res + "\t" + cpg.dev + "\n");
		}
		obuf.close();
	}

	// helper class to parse the config file
	public static class Config {
		int num_rep = 0;
		Map<String, List<Replicate>> expt_tab = new LinkedHashMap<String, List<Replicate>>();

		// constructor
		public Config(String config) throws Exception {
			BufferedReader obuf = new BufferedReader(new FileReader(config));
			String line = obuf.readLine();
			while ((line = obuf.readLine()) != null) {
				// check for blank lines
				if (line.length() == 0) {
					continue;
				}

				// parse the line
				String arr[] = line.split("\t");
				String expt_cond = arr[0];
				List<Replicate> rep_lst = expt_tab.get(expt_cond);
				if (rep_lst == null) {
					rep_lst = new ArrayList<Replicate>();
					expt_tab.put(expt_cond, rep_lst);
				}
				rep_lst.add(new Replicate(arr[1]));
				num_rep++;
			}
			obuf.close();
		}
	}

	// class to store data for replicates
	public static class Replicate {
		String file;
		static final int CHR_COL = 0, POS_COL = 1, METH_READ_COL = 2,
				NOT_METH_READ_COL = 3;

		// constructor
		public Replicate(String file) {
			this.file = file;
		}
	}

	// helper class for doing the merging of the data files
	public static class DataPoint {
		int meth_read, not_meth_read;

		// constructor
		public DataPoint(int meth_read, int not_meth_read) {
			this.meth_read = meth_read;
			this.not_meth_read = not_meth_read;
		}
	}

	// CpG class to sort the data
	public static class CpG implements Comparable<CpG> {
		String chr;
		int pos;
		double dev = DEFAULT_INT;

		// constructor
		public CpG(String chr, int pos) {
			this.chr = chr;
			this.pos = pos;
		}

		// compare with some other CpG
		public int compareTo(CpG that) {
			int res = this.chr.compareTo(that.chr);

			// check for same chromosome
			if (res == 0) {
				if (this.pos < that.pos) {
					res = -1;
				} else if (this.pos == that.pos) {
					res = 0;
				} else {
					res = 1;
				}
			}
			return res;
		}

		// hash code
		public int hashCode() {
			return toString().hashCode();
		}

		// to string
		public String toString() {
			return chr + ":" + pos;
		}
	}

	// class to manage the p-values
	public static class PValue implements Comparable<PValue> {
		double dev, pval;

		// constructor
		public PValue(double dev, double pval) {
			this.dev = dev;
			this.pval = pval;
		}

		// comparable [lower dev => higher pval]
		public int compareTo(PValue that) {
			int res = 0;
			if (this.dev < that.dev) {
				res = 1;
			} else {
				res = -1;
			}
			return res;
		}
	}
}
